from matplotlib import pyplot

from bfitter import BinnedFitter
from fitter import Fitter


class FitterBreitWigner(Fitter):
    fcn_latex = r"$\frac{A}{(x^2 - m^2)^2 + m^2\Gamma^2}$"
    fcn_str = "{0}/((x**2 - {1}**2)**2 + {1}**2*{2}**2)"
    
    def __init__(self, x, y, yerr, par_names=r'$A$ $m$ $\Gamma$'.split()):
        super(FitterBreitWigner, self).__init__(x, y, par_names, yerr=yerr, fcn_str=self.fcn_str)

class BinnedFitterBreitWigner(BinnedFitter):
    fcn_latex = "$\\frac{A}{(x^2 - m^2)^2 + m^2\\Gamma^2}$"
    fcn_str = "{0}/((x**2 - {1}**2)**2 + {1}**2*{2}**2)"

    def __init__(self, data, par_names="A $m$ $\\Gamma$".split()):
        super(BinnedFitterBreitWigner, self).__init__(data, par_names)
        self._par_text_offset = 0.45
                                         
    def _set_initpars(self, frange):
        if not self.hist:
            raise FitException("BUG in bw.py. Need to set self.hist before calling _set_initpars.")

        bins = self.hist.bins()

        max_bin = max(bins)

        max_bin_index = 0
        for i in xrange(len(bins)):
            if bins[i] == max_bin:
                max_bin_index = i
                break
        
        return [0.08*max_bin, self.data_mean, self.data_std]


#class BinnedFitterBWBG(Fitter):
#    fcn_latex = "$\\frac{A}{(x^2 - m^2)^2 + m^2\\Gamma^2}$ \n\n $ + p_0 + p_1x + p_2x^2 + p_3x^3$"
#
#    def __init__(self, data, par_names='$A$ $m$ $\\Gamma$ $p_0$ $p_1$ $p_2$ $p_3$'.split()):
#        super(BinnedFitterBWBG, self).__init__(data, par_names)
#
#        self._par_text_offset = 0.35
#                                                             
#    def fcn(self, p, x):
#        return p[0]/((x**2 - p[1]**2)**2 + p[1]**2*p[2]**2) 
#    
#    def bg_fcn(self, p, x):
#        return p[3] + p[4]*x + p[5]*x**2 + p[6]*x**3
#
#    def _set_initpars(self, bins, bin_low_edges):
#        max_bin = max(bins)
#
#        max_bin_index = 0
#        for i in xrange(len(bins)):
#            if bins[i] == max_bin:
#                max_bin_index = i
#                break
#        
#        return [0.08*max_bin, self.data_mean, self.data_std, 0, 0, 0, 0]


